{
 "cells": [
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-01T05:38:14.244407Z",
     "start_time": "2025-06-01T05:38:10.973674Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 从modelscope下载\n",
    "from modelscope.msdatasets import MsDataset\n",
    "\n",
    "# 加载本地train.jsonl\n",
    "ms_dataset = MsDataset.load('./chatbot.jsonl')\n",
    "\n",
    "# MsDataset 转换为 HuggingFace 的 Dataset\n",
    "hf_dataset = ms_dataset.to_hf_dataset()\n",
    "\n",
    "hf_dataset[:5]"
   ],
   "id": "c2912e2519e3352d",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-06-01 13:38:12,423 - modelscope - WARNING - Use trust_remote_code=True. Will invoke codes from ./chatbot.jsonl. Please make sure that you can trust the external codes.\n",
      "2025-06-01 13:38:12,424 - modelscope - WARNING - Use trust_remote_code=True. Will invoke codes from json. Please make sure that you can trust the external codes.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'question': ['护士是什么', '建筑师是什么', '厨师是什么', '记者是什么', '会计师是什么'],\n",
       " 'answer': ['护士是在医院、诊所或其他医疗机构工作的医疗保健专业人士，他们提供病人护理、执行医嘱和协助医生进行治疗。',\n",
       "  '建筑师是设计建筑物和规划空间的专业人士，他们负责创造美观、实用和安全的建筑环境。',\n",
       "  '厨师是在餐馆、酒店或食品服务行业工作的烹饪专业人士，他们准备和制作食物，确保菜品的质量和味道。',\n",
       "  '记者是收集、核实和报道新闻的专业人士，他们通过采访、写作和编辑来传达信息，向公众提供新闻和故事。',\n",
       "  '会计师是管理和分析财务信息的专业人士，他们记录交易、编制财务报表并提供税务和财务管理建议。']}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-01T05:38:18.089807Z",
     "start_time": "2025-06-01T05:38:16.263700Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 从HuggingFace下载\n",
    "from transformers import GPT2Tokenizer\n",
    "\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")\n",
    "\n",
    "# 从ModelScope下载\n",
    "# from modelscope import GPT2Tokenizer, GPT2Model\n",
    "# tokenizer = GPT2Tokenizer.from_pretrained('openai-community/gpt2')"
   ],
   "id": "6b6a86037ed62719",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-01T05:38:20.214804Z",
     "start_time": "2025-06-01T05:38:20.209999Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch.utils.data import Dataset\n",
    "\n",
    "prompt = \"问题:{}\\n答案:\"\n",
    "\n",
    "class ChatDataset(Dataset):\n",
    "    def __init__(self, hf_dataset):\n",
    "        self.hf_dataset = hf_dataset\n",
    "        self.data = []\n",
    "        for qa in hf_dataset:  # 遍历文件的每一行\n",
    "            # 将问题+答案拼接为单序列\n",
    "            text = prompt.format(qa['question']) + qa['answer'] + tokenizer.eos_token\n",
    "            tokens = tokenizer.encode(text, return_tensors=\"pt\")[0]\n",
    "            self.data.append(tokens)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx]\n"
   ],
   "id": "26865359b045d785",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-01T05:38:24.202930Z",
     "start_time": "2025-06-01T05:38:24.184965Z"
    }
   },
   "cell_type": "code",
   "source": [
    "chat_dataset = ChatDataset(hf_dataset)\n",
    "chat_dataset[:1]"
   ],
   "id": "3d2aee70181485cf",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([29785,   106,   165,    95,   246,    25,   162,   232,    97, 18803,\n",
       "         42468, 20015,   222, 20046,   230,   198,   163, 18433,   162,    94,\n",
       "           230,    25,   162,   232,    97, 18803, 42468, 28839,   101, 44293,\n",
       "           119,   165,   247,    95, 23513, 46237,   232, 33699,   222, 22755,\n",
       "           244, 17739,   114, 20015,   244, 44293,   119,   163,   244,   245,\n",
       "         17312,   118,   162,   252,   226, 32432,    98, 43291, 21410, 44293,\n",
       "           119,   163,   244,   245, 46479,   251,   161,   223,    98, 10310,\n",
       "           241, 10310,   248, 21689, 18803,   171,   120,   234, 20015,   244,\n",
       "         20015,   105,   162,   237,   238,   160,   122,   249,   163,   245,\n",
       "           227, 21689,   162,   232,    97, 49426,   228, 23513, 33699,   100,\n",
       "         26193,   234, 44293,   119,   161,   246,   109,   161,   240,   234,\n",
       "         39355,   237, 27950,   102, 44293,   119, 37955, 32573,   249, 26193,\n",
       "           234,   162,   110,   119,   163,   244,   245, 16764, 50256])]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-01T05:38:26.816103Z",
     "start_time": "2025-06-01T05:38:26.811526Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch.nn.utils.rnn import pad_sequence\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "def collate_fn(batch):\n",
    "    return pad_sequence(batch, batch_first=True, padding_value=tokenizer.eos_token_id)\n",
    "\n",
    "data_loader = DataLoader(chat_dataset, batch_size=1, shuffle=True, collate_fn=collate_fn)"
   ],
   "id": "63e184b15dee6c75",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-01T05:38:29.331592Z",
     "start_time": "2025-06-01T05:38:29.325326Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for batch in data_loader:\n",
    "    print(batch)\n",
    "    print(tokenizer.decode(batch[0]))\n",
    "    break"
   ],
   "id": "bcb7828f93320457",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[29785,   106,   165,    95,   246,    25,   161,   119,   118,   163,\n",
      "           255,   239, 30585,   230, 42468, 20015,   222, 20046,   230,   198,\n",
      "           163, 18433,   162,    94,   230,    25,   161,   119,   118,   163,\n",
      "           255,   239, 30585,   230, 42468,   164,   106,   122,   164,   106,\n",
      "            94,   161,   119,   118,   163,   255,   239, 31965,   102,   161,\n",
      "           240,   234,   164,   100,   226,   161,  7134,   163,   102,   118,\n",
      "         29785,   112, 21410, 10310,   241, 10310,   248, 21689, 18803,   171,\n",
      "           120,   234, 20015,   244, 20015,   105,   164,   112,   253,   164,\n",
      "           112,    96, 26344,   249, 34460,   254,   163,   122,   236,   164,\n",
      "           100,   224, 23513, 22522,   252, 18796,   101,   161,   240,   234,\n",
      "         22522,   231, 17739,   101, 21410,   161,   119,   118,   163,   255,\n",
      "           239,   163,   236,   107,   161,    95,   225, 16764, 50256]])\n",
      "问题:建筑师是什么\n",
      "答案:建筑师是设计建筑物和规划空间的专业人士，他们负责创造美观、实用和安全的建筑环境。<|endoftext|>\n"
     ]
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-01T05:40:02.766827Z",
     "start_time": "2025-06-01T05:38:49.094041Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from transformers import GPT2LMHeadModel, GPT2Config\n",
    "import torch\n",
    "\n",
    "NUM_EPOCHS = 50\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"mps:0\")\n",
    "\n",
    "config = GPT2Config()\n",
    "model = GPT2LMHeadModel(config)\n",
    "model.to(device)\n",
    "model.train()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "\n",
    "    total_loss = 0\n",
    "    step = 0\n",
    "    for batch in data_loader:\n",
    "\n",
    "        batch = batch.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(input_ids=batch, labels=batch)\n",
    "        loss = outputs.loss\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        # 每50个step打印一次损失\n",
    "        if step % 50 == 0:\n",
    "            print(f'Batch [{step}/{len(data_loader)}], Loss: {loss.item():.4f}')\n",
    "        step += 1\n",
    "\n",
    "    avg_loss = total_loss / len(data_loader)\n",
    "    print(f'Epoch [{epoch + 1}/{NUM_EPOCHS}], Loss: {avg_loss:.4f}')"
   ],
   "id": "66f44477426b2348",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch [0/11], Loss: 10.9509\n",
      "Epoch [1/50], Loss: 9.9242\n",
      "Batch [0/11], Loss: 8.8775\n",
      "Epoch [2/50], Loss: 8.8936\n",
      "Batch [0/11], Loss: 8.7891\n",
      "Epoch [3/50], Loss: 8.3187\n",
      "Batch [0/11], Loss: 7.9937\n",
      "Epoch [4/50], Loss: 7.5415\n",
      "Batch [0/11], Loss: 7.2013\n",
      "Epoch [5/50], Loss: 6.7189\n",
      "Batch [0/11], Loss: 6.3816\n",
      "Epoch [6/50], Loss: 6.0197\n",
      "Batch [0/11], Loss: 5.4715\n",
      "Epoch [7/50], Loss: 5.4463\n",
      "Batch [0/11], Loss: 5.3393\n",
      "Epoch [8/50], Loss: 4.9498\n",
      "Batch [0/11], Loss: 4.0726\n",
      "Epoch [9/50], Loss: 4.5543\n",
      "Batch [0/11], Loss: 4.2251\n",
      "Epoch [10/50], Loss: 4.1502\n",
      "Batch [0/11], Loss: 3.4724\n",
      "Epoch [11/50], Loss: 3.8151\n",
      "Batch [0/11], Loss: 3.8913\n",
      "Epoch [12/50], Loss: 3.5687\n",
      "Batch [0/11], Loss: 3.5081\n",
      "Epoch [13/50], Loss: 3.3126\n",
      "Batch [0/11], Loss: 2.7378\n",
      "Epoch [14/50], Loss: 3.0434\n",
      "Batch [0/11], Loss: 3.0289\n",
      "Epoch [15/50], Loss: 2.8220\n",
      "Batch [0/11], Loss: 2.1313\n",
      "Epoch [16/50], Loss: 2.6045\n",
      "Batch [0/11], Loss: 2.8083\n",
      "Epoch [17/50], Loss: 2.4478\n",
      "Batch [0/11], Loss: 2.0901\n",
      "Epoch [18/50], Loss: 2.2731\n",
      "Batch [0/11], Loss: 2.3148\n",
      "Epoch [19/50], Loss: 2.0764\n",
      "Batch [0/11], Loss: 2.1111\n",
      "Epoch [20/50], Loss: 1.9248\n",
      "Batch [0/11], Loss: 1.6427\n",
      "Epoch [21/50], Loss: 1.7887\n",
      "Batch [0/11], Loss: 1.5943\n",
      "Epoch [22/50], Loss: 1.6610\n",
      "Batch [0/11], Loss: 1.6680\n",
      "Epoch [23/50], Loss: 1.5231\n",
      "Batch [0/11], Loss: 1.1889\n",
      "Epoch [24/50], Loss: 1.4106\n",
      "Batch [0/11], Loss: 1.5094\n",
      "Epoch [25/50], Loss: 1.3152\n",
      "Batch [0/11], Loss: 0.9678\n",
      "Epoch [26/50], Loss: 1.2121\n",
      "Batch [0/11], Loss: 1.2902\n",
      "Epoch [27/50], Loss: 1.1242\n",
      "Batch [0/11], Loss: 1.3179\n",
      "Epoch [28/50], Loss: 1.0186\n",
      "Batch [0/11], Loss: 1.0940\n",
      "Epoch [29/50], Loss: 0.9455\n",
      "Batch [0/11], Loss: 0.9861\n",
      "Epoch [30/50], Loss: 0.8887\n",
      "Batch [0/11], Loss: 0.8341\n",
      "Epoch [31/50], Loss: 0.8162\n",
      "Batch [0/11], Loss: 0.6061\n",
      "Epoch [32/50], Loss: 0.7545\n",
      "Batch [0/11], Loss: 0.5386\n",
      "Epoch [33/50], Loss: 0.7143\n",
      "Batch [0/11], Loss: 0.8096\n",
      "Epoch [34/50], Loss: 0.6612\n",
      "Batch [0/11], Loss: 0.4777\n",
      "Epoch [35/50], Loss: 0.6031\n",
      "Batch [0/11], Loss: 0.4768\n",
      "Epoch [36/50], Loss: 0.5683\n",
      "Batch [0/11], Loss: 0.4578\n",
      "Epoch [37/50], Loss: 0.5346\n",
      "Batch [0/11], Loss: 0.5329\n",
      "Epoch [38/50], Loss: 0.4941\n",
      "Batch [0/11], Loss: 0.5117\n",
      "Epoch [39/50], Loss: 0.4667\n",
      "Batch [0/11], Loss: 0.3628\n",
      "Epoch [40/50], Loss: 0.4333\n",
      "Batch [0/11], Loss: 0.4493\n",
      "Epoch [41/50], Loss: 0.4021\n",
      "Batch [0/11], Loss: 0.3990\n",
      "Epoch [42/50], Loss: 0.3789\n",
      "Batch [0/11], Loss: 0.3724\n",
      "Epoch [43/50], Loss: 0.3601\n",
      "Batch [0/11], Loss: 0.3490\n",
      "Epoch [44/50], Loss: 0.3490\n",
      "Batch [0/11], Loss: 0.2928\n",
      "Epoch [45/50], Loss: 0.3233\n",
      "Batch [0/11], Loss: 0.3194\n",
      "Epoch [46/50], Loss: 0.3079\n",
      "Batch [0/11], Loss: 0.2975\n",
      "Epoch [47/50], Loss: 0.2927\n",
      "Batch [0/11], Loss: 0.2663\n",
      "Epoch [48/50], Loss: 0.2679\n",
      "Batch [0/11], Loss: 0.2552\n",
      "Epoch [49/50], Loss: 0.2660\n",
      "Batch [0/11], Loss: 0.2689\n",
      "Epoch [50/50], Loss: 0.2519\n"
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-01T14:09:47.680752Z",
     "start_time": "2025-06-01T14:09:47.670758Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 打印模型参数量\n",
    "print(f'参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad)}')"
   ],
   "id": "8921126c594aa12e",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "参数量：124439808\n"
     ]
    }
   ],
   "execution_count": 10
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-01T05:40:22.470795Z",
     "start_time": "2025-06-01T05:40:09.804901Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 设置生成参数\n",
    "generation_config = {\n",
    "    \"max_length\": 500,  # 生成最大长度（包括输入）\n",
    "    \"eos_token_id\": tokenizer.eos_token_id,  # 终止条件\n",
    "    \"pad_token_id\": tokenizer.eos_token_id,  # 若需填充，使用EOS的ID\n",
    "    \"do_sample\": True,  # 启用采样\n",
    "    \"temperature\": 0.8,  # 平衡确定性与随机性\n",
    "    \"top_k\": 50,  # 限制候选token数量\n",
    "    \"num_return_sequences\": 1,  # 生成1个不同结果\n",
    "}\n",
    "\n",
    "# 输入提示（可为空或部分诗句）\n",
    "input = prompt.format(\"程序员是什么\")\n",
    "input_ids = tokenizer.encode(input, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# 生成文本\n",
    "model.eval()\n",
    "outputs = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    **generation_config\n",
    ")\n",
    "\n",
    "# 解码并打印结果\n",
    "for i, output in enumerate(outputs):\n",
    "    result = tokenizer.decode(output, skip_special_tokens=True)\n",
    "    result = result[len(input):]\n",
    "    print(result)"
   ],
   "id": "df6f0e1df51b6b04",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "程序员是编写计算机软件代码的专业人士，他们使用编程语言实现算法和逻辑，开发应用程序和系统。\n"
     ]
    }
   ],
   "execution_count": 9
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
